#!/usr/bin/env python2
# -*- coding: utf-8 -*-
"""
sudo ./build/triton ./src/examples/pin/learn/taint.py  ./src/samples/crackmes/crackme_xor elite
read:0x7fff202b6626, size:1
read:0x7fff202b6627, size:1
read:0x7fff202b6628, size:1
read:0x7fff202b6629, size:1
read:0x7fff202b662a, size:1
Win
"""


from triton import ARCH, MemoryAccess, OPERAND
from pintool import *

Triton = getTritonContext()
def cbeforeSymProc(instruction):
    if instruction.getAddress() == 0x400556:
        rdi = getCurrentRegisterValue(Triton.registers.rdi)
        # 内存要对齐
        Triton.taintMemory(MemoryAccess(rdi, 8))


def cafter(inst):

    if inst.isTainted():
        # print('[tainted] %s' % (str(inst)))

        if inst.isMemoryRead():
            for op in inst.getOperands():
                if op.getType() == OPERAND.MEM:
                    print("read:0x{:08x}, size:{}".format(
                        op.getAddress(), op.getSize()))

        if inst.isMemoryWrite():
            for op in inst.getOperands():
                if op.getType() == OPERAND.MEM:
                    print("write:0x{:08x}, size:{}".format(
                        op.getAddress(), op.getSize()))


if __name__ == '__main__':
    startAnalysisFromSymbol('check')
    insertCall(cbeforeSymProc,  INSERT_POINT.BEFORE_SYMPROC)
    insertCall(cafter,          INSERT_POINT.AFTER)
    runProgram()
